'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 1_study_summary_memory.py
* @Time: 2025/7/20
* @All Rights Reserve By Brtc
'''
from typing import Any
import dotenv
from openai import OpenAI
"""
# 需求:
1、当里历史记录tokens 数量超国max_tokens 时则触发摘要总结
2、summary成员变量用于存储摘要
3、chat_histories用于存储历史对话
4、get_num_tokens用于计算对话的token 数量
5、save_context用于存储新一轮对话到历史对话中
6、get_buffer_string 用于将历史对话转换成字符串
7、load_memory_variables 用于加载记忆变量信息
8、summary_text用于将旧的摘要和传入的对话生成新的摘要
"""
class ConversationSummaryBufferMemory:
    def __init__(self, summary ='', chat_histories:list=None, max_tokens:int =300):
        self.m_summary = summary
        self.m_chat_histories = chat_histories
        self.m_max_tokens = max_tokens
        self.m_client = OpenAI(base_url="https://api.ephone.chat/v1")

    def summary_text(self, old_summary:str, new_line: str)-> str:
        """用于将旧的摘要和传入的对话生成新的摘要"""

        """ 新摘要生成的提示词 """
        prompt = f""" 
        你是一个强大的聊天机器人，请根据用户提供的谈话内容，总结摘要，并将其添加到先前提供的摘要中，返回一个新的摘要，除了新摘要其他任何数据都不要生成，如果用户的对话信息里有一些关键的信息，比方说姓名、爱好、性别、重要事件等等，这些全部都要包括在生成的摘要中，摘要尽可能要还原用户的对话记录。
        请不要将<example>标签里的数据当成实际的数据，这里的数据只是一个示例数据，告诉你该如何生成新摘要。
        <example>
        当前摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量。
        新的对话：
        Human：为什么你认为人工智能是一股向善的力量？
        AI：因为人工智能会帮助人类充分发挥潜力。
        新摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量，因为它将帮助人类充分发挥潜力。
        </example>
        =====================以下的数据是实际需要处理的数据=====================
        当前摘要：{old_summary}
        新的对话：
        {new_line}
        请帮用户将上面的信息生成新摘要。
        """

        """通过大模型去生成摘要"""
        completion = self.m_client.chat.completions.create(
            model = "gpt-3.5-turbo-16k",
            messages = [{"role":"user", "content":prompt}]
        )
        return completion.choices[0].message.content


    def get_buffer_string(self)->str:
        """获取所有聊天记录的字符串"""
        buffer : str = ""
        """通过遍历m_chat_histories 里面存储所有记录 组装成一个所有历史消息的 字符串 """
        for chat in self.m_chat_histories:
            buffer += f"Human:{chat.get('human')}\n AI:{chat.get('ai')}\n\n"
        return buffer


    def load_memory_variables(self)->dict[str, Any]:
        """用于加载记忆变量信息"""
        """获取所有的历史消息"""
        buffer_string = self.get_buffer_string()
        """把所有的历史消息组装到一个字典里面  键值是chat_histories， value 是所有历史消息"""
        return {
            "chat_histories":f"摘要:{self.m_summary}\n\n  历史信息:{buffer_string}"
        }


    @classmethod
    def get_num_tokens(cls, query:str)->int:
        """用于计算对话的token 数量"""
        """经验值 * 一般是字符串长度的1.5倍 需要精确计算请 使用tiktoken包"""
        return int(len(query)*1.5)

    def save_context(self, human_query:str, ai_content:str):
        """用于存储新一轮对话到历史对话中"""
        # 1、将历史消息以追加的方式存入到m_chat_histories列表中
        self.m_chat_histories.append({"human":human_query, "ai":ai_content})

        # 2、获取对话消息的token长度
        buffer_string = self.get_buffer_string()
        tokens = self.get_num_tokens(buffer_string)
        # 3、如果token 长度大于  设定的m_max_tokens， 那么则触发 摘要生成
        if tokens>self.m_max_tokens:
            #4、最远的一条历史记录就是 历史列表中的第一个
            first_chat = self.m_chat_histories[0]
            #5、把第一个的历史消息个过去摘要 合并成一个新摘要
            self.m_summary = self.summary_text(self.m_summary, f"Human:{first_chat.get('human')}\n AI:{first_chat.get('ai')}")
            #6、新摘要生成后删除最旧的一条摘要信息
            del self.m_chat_histories[0]


# 写一个简单的 聊一天
dotenv.load_dotenv()
chat_client = OpenAI(base_url="https://api.ephone.chat/v1")

if __name__ == "__main__":
    app_memory = ConversationSummaryBufferMemory("", [], 300)
    while True:
        # 1、人类输入问题
        human_query = input("Human:")

        if human_query=="exit":
            print("Goodbye Human!!")
            exit(0)
        # 2、 加载记忆
        memory_variables = app_memory.load_memory_variables()
        print(memory_variables)
        # 3、 将人类的问题 和  记忆合成一个新的 prompt
        anwser_prompt = (
            "你是一个强大的聊天机器人，请根据对应的用户上下文和用户的问题来回答！\n\n"
            f"{memory_variables.get("chat_histories")}\n"
            f"用户的提问是:{human_query}"
        )
        # 4、将带有人类问题 + 历史记忆的 prompt 询问 大模型
        completion = chat_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role":"system","content":"你是OpenAi研发的机器人，请根据用户的要求来回答问题！"},
                      {"role":"user",  "content":anwser_prompt}],
            stream=True
        )
        #5、 流式解析大模型回答的问题
        ai_content =""
        print("AI:", end="", flush=True)
        for chunk in completion:
            if len(chunk.choices) > 0:
                one_content = chunk.choices[0].delta.content
                if one_content is None:
                    break
                print(one_content, end="", flush=True)
                ai_content += one_content
        print("\n")
        #6、 将本轮对话的信息存储到 历史消息中
        app_memory.save_context(human_query, ai_content)



